import numpy
import torch
import torch.nn as nn
from torch import FloatTensor
import torchvision
from matplotlib import pyplot as plot


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.cnn = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=2),  # 输出 28*28*16
            nn.MaxPool2d(kernel_size=5, stride=2, padding=0),  # 输出 14*14*16
            nn.LeakyReLU(),
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1, padding=2),  # 输出 14*14*32
            nn.AvgPool2d(kernel_size=5, stride=2, padding=0),  # 输出 7*7*32
            nn.LeakyReLU()
        )
        self.dense = nn.Sequential(
            nn.Linear(512, 256),
            nn.LeakyReLU(),
            nn.Linear(256, 128),
            nn.BatchNorm1d(128),
            nn.LeakyReLU(),
            nn.Linear(128, 10),
            nn.LeakyReLU()
        )

    def forward(self, x):
        cnn_out = self.cnn(x)
        dense_input = cnn_out.view(cnn_out.size(0), -1)
        dense_out = self.dense(dense_input)
        return dense_out


def main():
    draw_x = []
    draw_y = []

    dataset = torchvision.datasets.mnist.MNIST(root="./mnist", download=True)
    fcn = CNN()
    print(fcn)
    loss_function = nn.MSELoss()
    opt = torch.optim.Adam(fcn.parameters(), lr=1e-2)
    train_data = dataset.train_data.resize(60000, 1, 28, 28)
    train_label = dataset.train_labels

    size = 60000
    batch = 20
    page_size = int(size / batch)

    for epoch in range(100):
        for pos in range(page_size):

            start = pos * batch
            end = (pos + 1) * batch

            x = FloatTensor(train_data.numpy()[start:end, :, :, :])
            y = numpy.zeros(shape=(batch, 10))
            yr = train_label[start:end]
            y[numpy.arange(batch), yr] = 1

            y_ = fcn.forward(x)
            loss = loss_function(y_, FloatTensor(y))
            print("LOSS: " + str(loss.detach().numpy()))
            draw_y.append(loss.detach().numpy())
            draw_x.append(epoch * page_size + pos)

            plot.plot(draw_x, draw_y)
            plot.show()

            opt.zero_grad()
            loss.backward()
            opt.step()


if __name__ == "__main__":
    main()
